# 导入tushare
import tushare as ts
import pandas as pd
import os
import time
import random
import common.base_profile as pf

# 设置你的 Tushare token
ts.set_token(pf.token)

# 创建Tushare数据接口
pro = ts.pro_api()


def filter_options(basic_csv_path='../data/opt_etf_basic.csv'):
    """
    从期权基础信息CSV文件中过滤符合条件的期权合约
    
    过滤逻辑：
    1. 结算月是每年的3，6，9，12月（季度月份）
    2. 认沽期权（call_put字段的值是 P）
    3. 一类认沽期权（opt_code，s_month相同）保留开始交易日期最早的合约
    4. 对过滤后的这一类合约，行权价格（exercise_price）排序，找到价格最低的合约
    
    参数:
    basic_csv_path (str): 期权基础信息CSV文件路径
    
    返回:
    pandas.DataFrame: 过滤后的期权合约列表
    """
    print(f"正在读取期权基础信息文件: {basic_csv_path}")
    df = pd.read_csv(basic_csv_path, encoding='utf-8-sig')
    print(f"原始数据共 {len(df)} 条记录")
    
    # 第一步：过滤结算月是3,6,9,12月的认沽期权
    # s_month格式是YYYYMM，提取月份部分
    df['month'] = df['s_month'].astype(str).str[-2:].astype(int)
    filtered_df = df[
        (df['month'].isin([3, 6, 9, 12])) &
        (df['call_put'] == 'P')
    ].copy()
    print(f"过滤后（结算月3,6,9,12月 + 认沽期权）共 {len(filtered_df)} 条记录")
    # INSERT_YOUR_CODE
    output_path = '../data/filtered_opt_etf_put_quarterly.csv'
    filtered_df.to_csv(output_path, index=False, encoding='utf-8-sig')
    print(f"已将过滤后的期权合约数据输出到: {output_path}")
    if filtered_df.empty:
        print("没有符合条件的期权合约")
        return pd.DataFrame()
    
    # 第二步和第三步合并：对于（opt_code，s_month）相同的合约
    # 先按list_date排序，找出最早的日期；在最早日期的合约中，选择exercise_price最低的
    def select_earliest_lowest_strike(group):
        """
        选择开始交易日期最早，且行权价格最低的合约
        先按list_date排序，再按exercise_price排序
        """
        group = group.sort_values(by=['list_date', 'exercise_price'])
        return group.iloc[0]
    
    final_df = filtered_df.groupby(['opt_code', 's_month']).apply(
        select_earliest_lowest_strike
    ).reset_index(drop=True)
    print(f"最终筛选后共 {len(final_df)} 条记录")
    # INSERT_YOUR_CODE
    # 结果输出到/data 目录里面
    output_path_data = '../data/opt_etf_quarterly_put_lowest_strike_daily.csv'
    final_df.to_csv(output_path_data, index=False, encoding='utf-8-sig')
    print(f"最终结果已输出到: {output_path_data}")
    return final_df


def fetch_opt_daily_data(ts_code, start_date=None, end_date=None):
    """
    获取期权日线数据
    
    参数:
    ts_code (str): 期权合约代码
    start_date (str): 开始日期，格式YYYYMMDD，可选
    end_date (str): 结束日期，格式YYYYMMDD，可选
    
    返回:
    pandas.DataFrame: 期权日线数据
    """
    try:
        # 添加随机延迟，避免请求过快
        time.sleep(0.2 + random.random() * 0.3)
        
        if start_date and end_date:
            df = pro.opt_daily(ts_code=ts_code, start_date=start_date, end_date=end_date)
        else:
            df = pro.opt_daily(ts_code=ts_code)
        
        return df
    except Exception as e:
        print(f"获取 {ts_code} 日线数据失败: {e}")
        return pd.DataFrame()


def save_daily_data(ts_code, daily_df, output_dir='../data/opt_etf_daily'):
    """
    保存期权日线数据到CSV文件
    
    参数:
    ts_code (str): 期权合约代码
    daily_df (pandas.DataFrame): 日线数据
    output_dir (str): 输出目录
    """
    if daily_df is None or daily_df.empty:
        print(f"  {ts_code} 没有数据可保存")
        return
    
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    
    # 生成文件名：使用ts_code作为文件名
    filename = f"{ts_code.replace('.', '_')}.csv"
    filepath = os.path.join(output_dir, filename)
    
    # 保存为CSV
    daily_df.to_csv(filepath, index=False, encoding='utf-8-sig')
    print(f"  {ts_code} 数据已保存到 {filepath}，共 {len(daily_df)} 条记录")


def main():
    """
    主函数：执行期权数据获取流程
    """
    print("=" * 60)
    print("开始获取期权日线数据")
    print("=" * 60)
    
    # 1. 过滤期权合约
    filtered_options = filter_options()
    
    if filtered_options.empty:
        print("没有符合条件的期权合约，程序退出")
        return
    
    print("\n符合条件的期权合约列表：")
    print(filtered_options[['ts_code', 'name', 'opt_code', 's_month', 'exercise_price', 'list_date']])
    print()
    
    # 2. 获取日线数据并保存
    output_dir = '../data/opt_etf_daily'
    os.makedirs(output_dir, exist_ok=True)
    
    success_count = 0
    fail_count = 0
    
    for idx, row in filtered_options.iterrows():
        ts_code = row['ts_code']
        name = row['name']
        opt_code = row['opt_code']
        s_month = row['s_month']
        exercise_price = row['exercise_price']
        
        print(f"[{idx + 1}/{len(filtered_options)}] 正在获取 {ts_code} ({name}) 的日线数据...")
        
        # 获取日线数据
        daily_df = fetch_opt_daily_data(ts_code)
        
        if daily_df is not None and not daily_df.empty:
            # 保存数据
            save_daily_data(ts_code, daily_df, output_dir)
            success_count += 1
        else:
            print(f"  {ts_code} 获取数据失败或返回空数据")
            fail_count += 1
    
    print("\n" + "=" * 60)
    print(f"数据获取完成！成功: {success_count} 个，失败: {fail_count} 个")
    print("=" * 60)


if __name__ == '__main__':
    main()

